import os
import sys
import importlib.util
import torch

# 获取当前工作目录
current_dir = os.path.dirname(os.path.abspath(__file__))

# 定义要搜索的根目录
model_root_dir = os.path.join(current_dir + "/" + "../", 'model')

# 用于存储导入的模块
imported_modules = {}
# print("current_dir",current_dir)
# print("model_root_dir",model_root_dir)
# 递归遍历目录
for root, dirs, files in os.walk(model_root_dir):
    # print("root",root)
    for file in files:
        # print("file",file)
        if file.endswith('.py'):  # 只处理 Python 文件
            file_path = os.path.join(root, file)
            # 生成唯一的模块名，使用相对路径来区分同名文件
            relative_path = os.path.relpath(file_path, model_root_dir)
            module_name = relative_path.replace(os.sep, '_').replace('.py', '')

            # 动态导入模块
            spec = importlib.util.spec_from_file_location(module_name, file_path)
            # print("spec",spec)
            if spec is not None:
                module = importlib.util.module_from_spec(spec)
                sys.modules[module_name] = module
                spec.loader.exec_module(module)

                # 将模块存储在字典中
                imported_modules[module_name] = module

# 测试导入的模块
# print(imported_modules['InDepthShareU_Net_InDepthShareU_Net'])

# for module_name, module in imported_modules.items():
# print(f"导入的模块: {module_name}")
# print(imported_modules)
# imported_modules['InDepthShareU_Net_InDepthShareU_Net'].U_Net(3,2)

# print("模块中的函数或类:", module)
# 这里可以调用模块中的函数或类，例如：

# if hasattr(module, 'UNet'):
#     tmp2_U_Nets.UNet(3,2)
# print("导入的模块数量:", len(imported_modules))


import NDSNet as nds
import U_Nets
import segNet
import UTNet.model.swin_unet as sw_unet
import patchNet
import SwinTransformer.models.swin_transformer as sw
#
# import InDepthShareU_Net.InDepthShareU_Net as tmp
#
#
#
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


#
#
#
# 跑不起来的模型
def NDSNet():
    return nds.NDSNet(3, 2).to(device)


# SwinUnet
def SwinUnet():
    return sw_unet.SwinUnet().to(device)


# PSPNet
def PSPNet():
    return patchNet.PSPNet().to(device)


# -------------------------------------------------------------------------

# 能正常运行的模型
# Unet
def UNet():
    return U_Nets.U_Net(3, 2).to(device)


# Unet++
def NestedUNet():
    return U_Nets.NestedUNet(3, 2).to(device)


# SegNet
def SegNet():
    return segNet.SegNet(3, 2).to(device)


# SWIN
def SwinTransformer():
    return sw.SwinTransformer().to(device)


def InDepthShareU_Net():
    return imported_modules['InDepthShareU_Net_InDepthShareU_Net'].U_Net(3, 2).to(device)


def tmp2():
    # return U_Nets.U_Net(3,2).to(device)
    return imported_modules['U_Nets_U_Nets'].U_Net(3, 2).to(device)


